#!/usr/bin/env python3
"""
关节状态发布器节点
以不同周期向四个话题发布消息：
- /hal/arm_joint_state: 10ms周期
- /hal/waist_state: 50ms周期  
- /mink_ik/target_arm_joints: 10ms周期
- /mink_ik/target_waist_joints: 10ms周期
"""

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_msgs.msg import Header
import numpy as np
import time

# Import genie_msgs for WaistState
from genie_msgs.msg import WaistState, MotorState

class JointStatePublisher(Node):
    """
    ROS2节点，向四个话题发布消息，使用不同的发布周期
    """
    
    def __init__(self):
        # 使用名称'joint_state_publisher'初始化节点
        super().__init__('joint_state_publisher')
        
        # 创建四个话题的发布者
        self.hal_arm_publisher = self.create_publisher(
            JointState, 
            '/hal/arm_joint_state', 
            10
        )
        
        self.hal_waist_publisher = self.create_publisher(
            WaistState, 
            '/hal/waist_state', 
            10
        )
        
        self.mink_ik_arm_publisher = self.create_publisher(
            JointState, 
            '/mink_ik/target_arm_joints', 
            10
        )
        
        self.mink_ik_waist_publisher = self.create_publisher(
            JointState, 
            '/mink_ik/target_waist_joints', 
            10
        )
        
        # 创建两个定时器，使用不同的周期
        # 10ms定时器：用于arm_joint_state和mink_ik话题
        self.fast_timer = self.create_timer(0.01, self.fast_timer_callback)
        
        # 20ms定时器：专门用于/hal/waist_state话题，提高频率避免冲突
        self.slow_timer = self.create_timer(0.02, self.slow_timer_callback)
        
        # 初始化关节名称
        self.arm_joint_names = [
            'arm_joint_1', 'arm_joint_2', 'arm_joint_3', 'arm_joint_4',
            'arm_joint_5', 'arm_joint_6', 'arm_joint_7', 'arm_joint_8',
            'arm_joint_9', 'arm_joint_10', 'arm_joint_11', 'arm_joint_12',
            'arm_joint_13', 'arm_joint_14'
        ]
        
        self.waist_joint_names = [
            'waist_pitch', 'lift'
        ]
        
        # 初始化关节位置
        self.arm_positions = [0.0] * 14
        # 将腰部初始位置设置在目标范围的中点
        # waist_pitch 范围 [0.429, 0.441] -> 中点 0.435
        # lift 范围 [0.29, 0.481] -> 中点 0.3855
        self.waist_positions = [0.435, 0.3855]
        
        # 计数器
        self.fast_counter = 0
        self.slow_counter = 0
        
        self.get_logger().info('关节状态发布器节点已启动')
        self.get_logger().info('发布周期配置:')
        self.get_logger().info('  - /hal/arm_joint_state: 100Hz (10ms)')
        self.get_logger().info('  - /hal/waist_state: 20Hz (50ms)')
        self.get_logger().info('  - /mink_ik/target_arm_joints: 100Hz (10ms)')
        self.get_logger().info('  - /mink_ik/target_waist_joints: 100Hz (10ms)')
    
    def fast_timer_callback(self):
        """
        快速定时器回调函数，每10ms被调用一次
        用于发布arm_joint_state和mink_ik话题
        """
        # 获取当前时间戳
        current_time = self.get_clock().now()
        current_time_msg = current_time.to_msg()
        
        # 1. 发布到 /hal/arm_joint_state (JointState)
        hal_arm_msg = JointState()
        hal_arm_msg.header.stamp = current_time_msg
        hal_arm_msg.name = self.arm_joint_names
        hal_arm_msg.position = self.arm_positions
        hal_arm_msg.velocity = [0.0] * 14
        hal_arm_msg.effort = [0.0] * 14
        self.hal_arm_publisher.publish(hal_arm_msg)
        
        # 2. 发布到 /mink_ik/target_arm_joints (JointState)
        mink_ik_arm_msg = JointState()
        mink_ik_arm_msg.header.stamp = current_time_msg
        mink_ik_arm_msg.name = self.arm_joint_names
        mink_ik_arm_msg.position = self.arm_positions
        mink_ik_arm_msg.velocity = [0.0] * 14
        mink_ik_arm_msg.effort = [0.0] * 14
        self.mink_ik_arm_publisher.publish(mink_ik_arm_msg)
        
        # 3. 发布到 /mink_ik/target_waist_joints (JointState)
        mink_ik_waist_msg = JointState()
        mink_ik_waist_msg.header.stamp = current_time_msg
        mink_ik_waist_msg.name = self.waist_joint_names
        mink_ik_waist_msg.position = self.waist_positions
        mink_ik_waist_msg.velocity = [0.0] * 2
        mink_ik_waist_msg.effort = [0.0] * 2
        self.mink_ik_waist_publisher.publish(mink_ik_waist_msg)
        
        # 更新关节位置用于演示
        self.update_joint_positions()
        
        # 每100次消息记录一次日志
        if self.fast_counter % 100 == 0:
            self.get_logger().info(f'快速定时器 - 已发布3个话题消息 - 计数器: {self.fast_counter}')
        
        self.fast_counter += 1
    
    def slow_timer_callback(self):
        """
        慢速定时器回调函数，每50ms被调用一次
        专门用于发布/hal/waist_state话题
        """
        # 获取当前时间戳
        current_time = self.get_clock().now()
        current_time_msg = current_time.to_msg()
        
        # 发布到 /hal/waist_state (WaistState) - 使用正确的genie_msgs类型
        hal_waist_msg = WaistState()
        hal_waist_msg.header.stamp = current_time_msg
        hal_waist_msg.name = self.waist_joint_names
        
        # 创建MotorState数组
        motor_states = []
        for i in range(2):
            motor_state = MotorState()
            motor_state.id = i
            motor_state.enable = True
            motor_state.position = self.waist_positions[i]
            motor_state.velocity = 0.0
            motor_state.effort = 0.0
            motor_state.current = 0.0
            motor_state.voltage = 12.0
            motor_state.temperature = 25.0
            motor_state.status = 0  # 正常状态
            motor_state.err_code = 0  # 无错误
            motor_states.append(motor_state)
        
        hal_waist_msg.motor_states = motor_states
        self.hal_waist_publisher.publish(hal_waist_msg)
        
        # 每20次消息记录一次日志（因为50ms周期，20次就是1秒）
        if self.slow_counter % 20 == 0:
            self.get_logger().info(f'慢速定时器 - 已发布/hal/waist_state消息 - 计数器: {self.slow_counter}')
        
        self.slow_counter += 1
    
    def update_joint_positions(self):
        """
        更新关节位置用于演示目的
        你可以修改此函数来生成不同的轨迹
        """
        # 示例：简单的正弦运动
        time_val = time.time()
        
        # 更新手臂关节位置
        for i in range(14):
            self.arm_positions[i] = 0.5 * np.sin(time_val + i * 0.5)
        
        # 更新腰部关节位置，严格限制在给定范围内
        # waist_pitch 范围 [0.429, 0.441]
        waist_pitch_min = 0.429
        waist_pitch_max = 0.441
        waist_pitch_center = (waist_pitch_min + waist_pitch_max) / 2.0
        waist_pitch_amp = (waist_pitch_max - waist_pitch_min) / 2.0
        self.waist_positions[0] = waist_pitch_center + waist_pitch_amp * np.sin(time_val)

        # lift 范围 [0.29, 0.481]
        lift_min = 0.29
        lift_max = 0.481
        lift_center = (lift_min + lift_max) / 2.0
        lift_amp = (lift_max - lift_min) / 2.0
        self.waist_positions[1] = lift_center + lift_amp * np.sin(time_val * 0.8 + 1.0)

def main(args=None):
    """
    主函数，用于初始化和运行ROS2节点
    """
    # 初始化ROS2 Python客户端库
    rclpy.init(args=args)
    
    # 创建节点
    joint_state_publisher = JointStatePublisher()
    
    try:
        # 旋转节点，使回调函数被调用
        rclpy.spin(joint_state_publisher)
    except KeyboardInterrupt:
        # 优雅地处理Ctrl+C
        pass
    finally:
        # 清理
        joint_state_publisher.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()